<template>
  <canvas ref="container" class="w-full h-full"></canvas>
</template>

<script setup>
import { ref, onMounted } from "vue";
import { WebGPU } from "../../lib/webgpu/WebGPU";
import {
  getModelViewMatrix,
  getProjectionMatrix,
} from "../../lib/math/GetMatrix";

const container = ref(null);
const webgpu = new WebGPU();
let oldTime = Date.now();

const cube = {
  vertex: new Float32Array([
    // float3 position, float3 normal
    // face1
    1, 1, 1, 1, 0, 0, 1, 1, -1, 1, 0, 0, 1, -1, 1, 1, 0, 0, 1, -1, -1, 1, 0, 0,
    // face2
    -1, 1, -1, -1, 0, 0, -1, 1, 1, -1, 0, 0, -1, -1, -1, -1, 0, 0, -1, -1, 1,
    -1, 0, 0,
    // face3
    -1, 1, -1, 0, 1, 0, 1, 1, -1, 0, 1, 0, -1, 1, 1, 0, 1, 0, 1, 1, 1, 0, 1, 0,
    // face4
    -1, -1, 1, 0, -1, 0, 1, -1, 1, 0, -1, 0, -1, -1, -1, 0, -1, 0, 1, -1, -1, 0,
    -1, 0,
    // face5
    -1, 1, 1, 0, 0, 1, 1, 1, 1, 0, 0, 1, -1, -1, 1, 0, 0, 1, 1, -1, 1, 0, 0, 1,
    // face6
    1, 1, -1, 0, 0, -1, -1, 1, -1, 0, 0, -1, 1, -1, -1, 0, 0, -1, -1, -1, -1, 0,
    0, -1,
  ]),
  index: new Uint16Array([
    0, 2, 1, 2, 3, 1, 4, 6, 5, 6, 7, 5, 8, 10, 9, 10, 11, 9, 12, 14, 13, 14, 15,
    13, 16, 18, 17, 18, 19, 17, 20, 22, 21, 22, 23, 21,
  ]),
  vertexCount: 24,
  indexCount: 36,
};

async function initPipeline(device, format, size) {
  let pipeline = await device.createRenderPipelineAsync({
    layout: "auto",
    vertex: {
      module: device.createShaderModule({
        code: `
          @binding(0) @group(0) var<uniform> modelViewMatrix : mat4x4<f32>;
          @binding(1) @group(0) var<uniform> projectionMatrix : mat4x4<f32>;

          struct VertexOutput {
              @builtin(position) Position : vec4<f32>,
              @location(0) fragPosition: vec3<f32>,
              @location(1) fragNormal : vec3<f32>,
          };

          @vertex
          fn main(
              @location(0) position : vec3<f32>,
              @location(1) normal : vec3<f32>
          ) -> VertexOutput {
              var mvpMatrix = projectionMatrix * modelViewMatrix;
              var pos = vec4<f32>(position, 1.0);
              var nor = vec4<f32>(normal, 0.0);

              var output : VertexOutput;
              output.Position = mvpMatrix * pos;
              output.fragPosition = (modelViewMatrix * pos).xyz;
              output.fragNormal = (modelViewMatrix * nor).xyz;
              return output;
          }
        `,
      }),
      entryPoint: "main",
      buffers: [
        {
          arrayStride: 6 * 4,
          attributes: [
            {
              shaderLocation: 0,
              offset: 0,
              format: "float32x3",
            },
            {
              shaderLocation: 1,
              offset: 3 * 4,
              format: "float32x3",
            },
          ],
        },
      ],
    },
    fragment: {
      module: device.createShaderModule({
        code: `
          @binding(0) @group(1) var<uniform> ambientLight : f32;
          @binding(1) @group(1) var<uniform> directionLight : array<vec4<f32>, 2>;
          @binding(2) @group(1) var<uniform> pointLight : array<vec4<f32>, 2>;
        
          @fragment
          fn main(
              @location(0) fragPosition: vec3<f32>,
              @location(1) fragNormal: vec3<f32>,
          ) -> @location(0) vec4<f32> {
              var color = vec3<f32>(1.0, 0.0, 0.0);
              var result = color * ambientLight;

              var directionPosition = directionLight[0].xyz;
              var directionIntensity = directionLight[1][0];
              var directionDiffuse = max(dot(normalize(directionPosition), fragNormal), 0.0);
              result += color * directionIntensity * directionDiffuse;

              var pointPosition = pointLight[0].xyz;
              var pointIntensity = pointLight[1][0];
              var pointRadius = pointLight[1][1];
              var pointDir = pointPosition - fragPosition;
              var distance = length(pointDir);
              if(distance < pointRadius){
                  var pointDiffuse = max(dot(normalize(pointDir), fragNormal), 0.0);
                  var distanceFactor = pow(1.0 - distance / pointRadius, 2.0);
                  result += color * pointIntensity * pointDiffuse * distanceFactor;
              }
              
              return vec4<f32>(result, 1.0);
          }
        `,
      }),
      entryPoint: "main",
      targets: [
        {
          format: format,
        },
      ],
    },
    primitive: {
      topology: "triangle-list",
      cullMode: "back",
      frontFace: "ccw",
    },
    depthStencil: {
      depthWriteEnabled: true,
      depthCompare: "less",
      format: "depth24plus",
    },
  });

  let depthTexture = device.createTexture({
    size,
    format: "depth24plus",
    // eslint-disable-next-line no-undef
    usage: GPUTextureUsage.RENDER_ATTACHMENT,
  });
  let depthView = depthTexture.createView();

  let vertexBuffer = device.createBuffer({
    label: "GPUBuffer store vertex",
    size: cube.vertex.byteLength,
    // eslint-disable-next-line no-undef
    usage: GPUBufferUsage.VERTEX | GPUBufferUsage.COPY_DST,
  });
  device.queue.writeBuffer(vertexBuffer, 0, cube.vertex);

  let indexBuffer = device.createBuffer({
    label: "GPUBuffer store vertex",
    size: cube.index.byteLength,
    // eslint-disable-next-line no-undef
    usage: GPUBufferUsage.INDEX | GPUBufferUsage.COPY_DST,
  });
  device.queue.writeBuffer(indexBuffer, 0, cube.index);

  let modelViewBuffer = device.createBuffer({
    label: "GPUBuffer store 4x4 matrix",
    size: 4 * 4 * 4,
    // eslint-disable-next-line no-undef
    usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
  });
  let projectionBuffer = device.createBuffer({
    label: "GPUBuffer store 4x4 matrix",
    size: 4 * 4 * 4,
    // eslint-disable-next-line no-undef
    usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
  });
  let vertexGroup = device.createBindGroup({
    label: "Uniform Group with Matrix",
    layout: pipeline.getBindGroupLayout(0),
    entries: [
      {
        binding: 0,
        resource: {
          buffer: modelViewBuffer,
        },
      },
      {
        binding: 1,
        resource: {
          buffer: projectionBuffer,
        },
      },
    ],
  });

  let ambientBuffer = device.createBuffer({
    label: "GPUBuffer store 4x4 matrix",
    size: 1 * 4,
    // eslint-disable-next-line no-undef
    usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
  });
  let directionalBuffer = device.createBuffer({
    label: "GPUBuffer store 4x4 matrix",
    size: 8 * 4,
    // eslint-disable-next-line no-undef
    usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
  });
  let pointBuffer = device.createBuffer({
    label: "GPUBuffer store 4x4 matrix",
    size: 8 * 4,
    // eslint-disable-next-line no-undef
    usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
  });
  let lightGroup = device.createBindGroup({
    label: "Uniform Group with Matrix",
    layout: pipeline.getBindGroupLayout(1),
    entries: [
      {
        binding: 0,
        resource: {
          buffer: ambientBuffer,
        },
      },
      {
        binding: 1,
        resource: {
          buffer: directionalBuffer,
        },
      },
      {
        binding: 2,
        resource: {
          buffer: pointBuffer,
        },
      },
    ],
  });

  return {
    pipeline,
    vertexBuffer,
    indexBuffer,
    depthTexture,
    depthView,
    vertexGroup,
    modelViewBuffer,
    projectionBuffer,
    lightGroup,
    ambientBuffer,
    directionalBuffer,
    pointBuffer,
  };
}

function draw(device, context, pipelineObj) {
  const commandEncoder = device.createCommandEncoder();
  const renderPassDescriptor = {
    colorAttachments: [
      {
        view: context.getCurrentTexture().createView(),
        clearValue: { r: 0, g: 0, b: 0, a: 1.0 },
        loadOp: "clear",
        storeOp: "store",
      },
    ],
    depthStencilAttachment: {
      view: pipelineObj.depthView,
      depthClearValue: 1.0,
      depthLoadOp: "clear",
      depthStoreOp: "store",
    },
  };
  const passEncoder = commandEncoder.beginRenderPass(renderPassDescriptor);
  passEncoder.setPipeline(pipelineObj.pipeline);
  passEncoder.setVertexBuffer(0, pipelineObj.vertexBuffer);
  passEncoder.setIndexBuffer(pipelineObj.indexBuffer, "uint16");
  passEncoder.setBindGroup(0, pipelineObj.vertexGroup);
  passEncoder.setBindGroup(1, pipelineObj.lightGroup);
  passEncoder.drawIndexed(36, 1, 0, 0, 0);
  passEncoder.end();
  device.queue.submit([commandEncoder.finish()]);
}

async function init() {
  let canvas = container.value;
  let { device, context, format, size } = await webgpu.initWebGPU(canvas);
  let pipelineObj = await initPipeline(device, format, size);

  let aspect = size.width / size.height;
  let position = { x: 0, y: 0, z: -5 };
  let rotation = { x: 0.5, y: 0, z: 0 };
  let scale = { x: 1, y: 1, z: 1 };

  let ambient = new Float32Array([0.1]);
  device.queue.writeBuffer(pipelineObj.ambientBuffer, 0, ambient);
  let direction = new Float32Array([1, 0, 1, 1, 0.5, 0, 0, 0]);
  device.queue.writeBuffer(pipelineObj.directionalBuffer, 0, direction);
  let point = new Float32Array([-5, 0, -5, 1, 1, 20, 0, 0]);
  device.queue.writeBuffer(pipelineObj.pointBuffer, 0, point);

  function frame() {
    let newTime = Date.now();
    let delta = (newTime - oldTime) / 1000;
    oldTime = newTime;

    rotation.y -= delta * 1;
    let modelViewMatrix = getModelViewMatrix(position, rotation, scale);
    device.queue.writeBuffer(
      pipelineObj.modelViewBuffer,
      0,
      modelViewMatrix.buffer
    );
    let projectionMatrix = getProjectionMatrix(aspect);
    device.queue.writeBuffer(
      pipelineObj.projectionBuffer,
      0,
      projectionMatrix.buffer
    );
    draw(device, context, pipelineObj);

    requestAnimationFrame(frame);
  }
  frame();

  window.addEventListener("resize", () => {
    size.width = canvas.width = canvas.clientWidth;
    size.height = canvas.height = canvas.clientHeight;
    pipelineObj.depthTexture.destroy();
    pipelineObj.depthTexture = device.createTexture({
      size,
      format: "depth24plus",
      // eslint-disable-next-line no-undef
      usage: GPUTextureUsage.RENDER_ATTACHMENT,
    });
    pipelineObj.depthView = pipelineObj.depthTexture.createView();
    aspect = size.width / size.height;
  });
}

onMounted(() => {
  init();
});
</script>
